# coding: utf-8
# create by tongshiwei on 2018/11/10
import random
from RecAgent.Agent import Agent

from longling.lib.utilog import config_logging, LogLevel
from collections import Counter


class CAgent(Agent):
    def __init__(self, graph, kt, model_dir, path_warm_up=True, window_size=None, visit_threshold=1,
                 allow_shortcut=True):

        super(CAgent, self).__init__(graph)

        self.kt = kt
        self.max_steps = 0
        self.model_dir = model_dir

        self._state = None
        self._target = None
        self.path = None
        self.pre_path = None
        self._legal_candidates = None

        self.window_size = window_size
        self.visit_threshold = visit_threshold
        self.allow_shortcut = allow_shortcut
        self.path_warm_up = path_warm_up

        self._cnt = 0

    @property
    def kt_state(self):
        return self._state[1]

    @property
    def mastery(self):
        return self._state[0]

    @property
    def state(self):
        return self.mastery

    def begin_episode(self, exercises_record, target, *args, **kwargs):
        mastery, _state = self.kt(exercises_record)
        self._state = (mastery[-1], _state)
        self._target = target if target is not None else list(self.graph.nodes)
        self.path = []
        self.pre_path = [e for e, c in exercises_record]
        self._legal_candidates = set(target)

    def test_correct(self, exercise, mastery):
        return 1 if 0.5 <= mastery[exercise] else 0

    def test(self, exercises, score_type=False, mastery=None):
        mastery = mastery if mastery is not None else self.mastery
        if score_type:
            return [(exercise, mastery[exercise]) for exercise in exercises]
        return [(exercise, self.test_correct(exercise, mastery)) for exercise in exercises]

    def test_score(self, exercises, score_type=False, mastery=None):
        return sum([s for _, s in self.test(exercises, score_type, mastery)])

    def search(self):
        return self.graph

    def step(self, *args, **kwargs):
        self.max_steps += 1

        candidates = self.get_candidate_actions()

        if candidates:
            return random.choice(candidates), None
        else:
            return None, None

    def state_transform(self, exercise, correct, add_to_path=True, *args, **kwargs):
        mastery, state = self.kt(
            [(exercise, self.test_correct(exercise, self.mastery) if correct is None else correct)], self.kt_state)
        self._state = (mastery[-1], state)

        if add_to_path is True:
            self.path.append(exercise)

    def end_episode(self):
        rec_path = self.search()
        self._state = None
        self._target = None
        self.path = None
        self.max_steps = 0
        self.pre_path = None
        self._legal_candidates = None
        return rec_path

    def get_candidate_actions(self):
        _candidates = set()
        _check_done = set()
        _check = list(self._target)

        while _check:
            node = _check.pop()
            if node in _check_done:
                continue
            if self.test_correct(node, self.mastery) == 1:
                _check_done.add(node)
                continue
            for pre_node in self.graph.predecessors(node):
                if self.test_correct(pre_node, self.mastery) != 1:
                    _check_done.add(node)
                    break
            else:
                _candidates.add(node)
                _check_done.add(node)

        return list(_candidates)


def main(dataset, ctx, max_steps, max_num, score_type, allow_shortcut=True, test_pre='', disable=False):
    import os
    from Graph import Graph
    from Environment import load_evaluator, load_dkt
    from mxnet import gpu
    from longling import wf_open

    model_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../data/%s/" % dataset)) + os.sep

    ctx = gpu(ctx)
    evaluator = load_evaluator(dataset, ctx, score_type,
                               logger=config_logging(logger="evaluator", console_log_level=LogLevel.INFO))
    kt = load_dkt(dataset, ctx)

    graph = Graph(dataset=dataset, disable=disable)
    agent = CAgent(
        graph=graph,
        kt=kt,
        model_dir=model_dir,
        path_warm_up=True,
        allow_shortcut=allow_shortcut,
    )
    method = 'Cog'
    if graph_disable is True:
        method += '-ng'
    else:
        if allow_shortcut is False:
            method += '-nsc'
    if score_type is True:
        method += '-score'

    root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../data/%s/" % dataset)) + os.sep
    model_dir = root_dir + "Baseline/%s/PL-%s/" % (method, max_steps)
    prefix = model_dir + "Cog%s-%s" % (test_pre, max_num)

    with wf_open(prefix + "_result.log") as wf, wf_open(prefix + "_path.log") as pwf:
        print(evaluator.eval(
            agent=agent,
            max_steps=max_steps,
            epsilon=0,
            alternative_mode=False,
            log_f=pwf,
            max_num=max_num,
            delay_estimation=False,
        ), file=wf)

    return evaluator


if __name__ == '__main__':
    import sys

    assert len(sys.argv) == 5

    dataset = sys.argv[1]
    ctx = int(sys.argv[2])
    max_steps = int(sys.argv[3])
    test_pre = sys.argv[4]

    # dataset = "irt"
    # ctx = 0
    # max_steps = 20
    # test_pre = "0"

    max_num = 1000
    score_type = False
    allow_shortcut = False
    graph_disable = False

    evaluator = main(
        dataset=dataset,
        ctx=ctx,
        max_steps=max_steps,
        max_num=max_num,
        score_type=score_type,
        allow_shortcut=allow_shortcut,
        test_pre=test_pre,
        disable=graph_disable,
    )
